{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "82f90261",
   "metadata": {},
   "source": [
    "# Kùzu Graph Store\n",
    "\n",
    "This notebook walks through configuring `Kùzu` to be the backend for graph storage in LlamaIndex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1a9eb90-335c-4214-8bb6-fd1edbe3ccbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My OpenAI Key\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"API_KEY_HERE\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a9f2e3-c729-455a-a338-2f83776c1d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "02299c66",
   "metadata": {},
   "source": [
    "## Prepare for Kùzu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d847733c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up all the directories used in this notebook\n",
    "import shutil\n",
    "\n",
    "shutil.rmtree(\"./test1\", ignore_errors=True)\n",
    "shutil.rmtree(\"./test2\", ignore_errors=True)\n",
    "shutil.rmtree(\"./test3\", ignore_errors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22aa9152",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting kuzu\n",
      "  Downloading kuzu-0.0.6-cp39-cp39-macosx_11_0_arm64.whl (5.5 MB)\n",
      "\u001b[K     |████████████████████████████████| 5.5 MB 4.8 MB/s eta 0:00:01\n",
      "\u001b[?25hInstalling collected packages: kuzu\n",
      "Successfully installed kuzu-0.0.6\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.2.1 is available.\n",
      "You should consider upgrading via the '/Users/loganmarkewich/llama_index/llama-index/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install kuzu\n",
    "import kuzu\n",
    "\n",
    "db = kuzu.Database(\"test1\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "be3f7baa-1c0a-430b-981b-83ddca9e71f2",
   "metadata": {},
   "source": [
    "## Using Knowledge Graph with KuzuGraphStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c8a77a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.graph_stores import KuzuGraphStore\n",
    "\n",
    "graph_store = KuzuGraphStore(db)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "75f1d565-04e8-41bc-9165-166dc89b6b47",
   "metadata": {},
   "source": [
    "#### Building the Knowledge Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d0b2364-4806-4656-81e7-3f6e4b910b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    SimpleDirectoryReader,\n",
    "    ServiceContext,\n",
    "    KnowledgeGraphIndex,\n",
    ")\n",
    "\n",
    "from llama_index.llms import OpenAI\n",
    "from IPython.display import Markdown, display\n",
    "import kuzu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c297fd3-3424-41d8-9d0d-25fe6310ab62",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\n",
    "    \"../../../../examples/paul_graham_essay/data\"\n",
    ").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61679142-7595-492b-8792-26cbc439caf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define LLM\n",
    "\n",
    "llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm, chunk_size=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "370fd08f-56ff-4c24-b0c4-c93116a6d482",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.storage.storage_context import StorageContext\n",
    "\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
    "\n",
    "# NOTE: can take a while!\n",
    "index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    max_triplets_per_chunk=2,\n",
    "    storage_context=storage_context,\n",
    "    service_context=service_context,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c39a0eeb-ef16-4982-8ba8-b37c2c5f4437",
   "metadata": {},
   "source": [
    "#### Querying the Knowledge Graph\n",
    "\n",
    "First, we can query and send only the triplets to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670300d8-d0a8-4201-bbcd-4a74b199fcdd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.knowledge_graph.retriever:> Starting query: Tell me more about Interleaf\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.knowledge_graph.retriever:> Query keywords: ['Interleaf']\n",
      "ERROR:llama_index.indices.knowledge_graph.retriever:Index was not constructed with embeddings, skipping embedding usage...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Extracted relationships: The following are knowledge sequence in max depth 2 in the form of `subject [predicate, object, predicate_next_hop, object_next_hop ...]`\n",
      "Interleaf ['made', 'software for creating documents']\n",
      "Interleaf ['added', 'scripting language']\n",
      "Interleaf ['taught', 'what not to do']\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=False, response_mode=\"tree_summarize\"\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about Interleaf\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eecf2d57-3efa-4b0d-941a-95438d42893c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>Interleaf is a company that made software for creating documents. They also added a scripting language to their software. Additionally, they taught what not to do.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9eeaba8a",
   "metadata": {},
   "source": [
    "For more detailed answers, we can also send the text from where the retrieved tripets were extracted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd14686d-1c53-4637-9340-3745f2121ae2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.knowledge_graph.retriever:> Starting query: Tell me more about Interleaf\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Query keywords: ['Interleaf']\n",
      "ERROR:llama_index.indices.knowledge_graph.retriever:Index was not constructed with embeddings, skipping embedding usage...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Querying with idx: 144f784c-d052-4fed-86f8-c895da6e13df: each student had. But the Accademia wasn't teaching me anything except Italia...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Querying with idx: 7c877dd3-3375-4ab7-8745-e0dfbabfe5bd: learned some useful things at Interleaf, though they were mostly about what n...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Extracted relationships: The following are knowledge sequence in max depth 2 in the form of `subject [predicate, object, predicate_next_hop, object_next_hop ...]`\n",
      "Interleaf ['made', 'software for creating documents']\n",
      "Interleaf ['added', 'scripting language']\n",
      "Interleaf ['taught', 'what not to do']\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=True, response_mode=\"tree_summarize\"\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about Interleaf\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c87d14-d2d8-4d80-89f6-1e5972973528",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>Interleaf was a company that made software for creating documents. They were inspired by Emacs and added a scripting language to their software, which was a dialect of Lisp. The company hired a Lisp hacker to write things in this scripting language. The narrator worked at Interleaf for a year but admits to being a bad employee. They found it difficult to understand most of the software because it was primarily written in C, a language they did not know or want to learn. Despite this, they were paid well and managed to save enough money to go back to RISD and pay off their college loans. The narrator also learned some valuable lessons at Interleaf, such as the importance of having product people rather than sales people running technology companies, the drawbacks of having too many people edit code, the impact of office space on productivity, the value of corridor conversations over planned meetings, the challenges of dealing with big bureaucratic customers, and the importance of being the \"entry level\" option in a market.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ecc7342a",
   "metadata": {},
   "source": [
    "#### Query with embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20f9da1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:llama_index.llms.openai_utils:Retrying llama_index.llms.openai_utils.completion_with_retry.<locals>._completion_with_retry in 4.0 seconds as it raised ServiceUnavailableError: The server is overloaded or not ready yet..\n"
     ]
    }
   ],
   "source": [
    "# NOTE: can take a while!\n",
    "db = kuzu.Database(\"test2\")\n",
    "graph_store = KuzuGraphStore(db)\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
    "new_index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    max_triplets_per_chunk=2,\n",
    "    service_context=service_context,\n",
    "    storage_context=storage_context,\n",
    "    include_embeddings=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1729841e",
   "metadata": {},
   "outputs": [],
   "source": [
    "rel_map = graph_store.get_rel_map()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01b74b2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.knowledge_graph.retriever:> Starting query: Tell me more about what the author worked on at Interleaf\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Query keywords: ['Interleaf', 'author', 'worked']\n",
      "ERROR:llama_index.indices.knowledge_graph.retriever:Index was not constructed with embeddings, skipping embedding usage...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Querying with idx: 144f784c-d052-4fed-86f8-c895da6e13df: each student had. But the Accademia wasn't teaching me anything except Italia...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Querying with idx: 7c877dd3-3375-4ab7-8745-e0dfbabfe5bd: learned some useful things at Interleaf, though they were mostly about what n...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Extracted relationships: The following are knowledge sequence in max depth 2 in the form of `subject [predicate, object, predicate_next_hop, object_next_hop ...]`\n",
      "Interleaf ['made', 'software for creating documents']\n",
      "Interleaf ['added', 'scripting language']\n",
      "Interleaf ['taught', 'what not to do']\n"
     ]
    }
   ],
   "source": [
    "# query using top 3 triplets plus keywords (duplicate triplets are removed)\n",
    "query_engine = index.as_query_engine(\n",
    "    include_text=True,\n",
    "    response_mode=\"tree_summarize\",\n",
    "    embedding_mode=\"hybrid\",\n",
    "    similarity_top_k=5,\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about what the author worked on at Interleaf\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02084f6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>At Interleaf, the author worked on creating software for creating documents. They also worked on adding a scripting language, which was inspired by Emacs and was a dialect of Lisp. However, the author admits to being a bad employee and not fully understanding the software, as it was primarily written in C. They also mention that they spent a lot of time working on their book \"On Lisp\" during their time at Interleaf. Overall, the author learned some useful things at Interleaf, particularly about what not to do in technology companies.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cd582500-584c-409a-9963-921738f1beb8",
   "metadata": {},
   "source": [
    "#### Visualizing the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1926b53e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting pyvis\n",
      "  Downloading pyvis-0.3.2-py3-none-any.whl (756 kB)\n",
      "\u001b[K     |████████████████████████████████| 756 kB 2.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jsonpickle>=1.4.1\n",
      "  Downloading jsonpickle-3.0.1-py2.py3-none-any.whl (40 kB)\n",
      "\u001b[K     |████████████████████████████████| 40 kB 4.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: networkx>=1.11 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pyvis) (3.1)\n",
      "Requirement already satisfied: ipython>=5.3.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pyvis) (8.10.0)\n",
      "Requirement already satisfied: jinja2>=2.9.6 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pyvis) (3.1.2)\n",
      "Requirement already satisfied: pexpect>4.3 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (4.8.0)\n",
      "Requirement already satisfied: backcall in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.2.0)\n",
      "Requirement already satisfied: decorator in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (5.1.1)\n",
      "Requirement already satisfied: pickleshare in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.7.5)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.30 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (3.0.39)\n",
      "Requirement already satisfied: appnope in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.1.3)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (2.15.1)\n",
      "Requirement already satisfied: traitlets>=5 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (5.9.0)\n",
      "Requirement already satisfied: jedi>=0.16 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.18.2)\n",
      "Requirement already satisfied: matplotlib-inline in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.1.6)\n",
      "Requirement already satisfied: stack-data in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.6.2)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from jedi>=0.16->ipython>=5.3.0->pyvis) (0.8.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from jinja2>=2.9.6->pyvis) (2.1.3)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pexpect>4.3->ipython>=5.3.0->pyvis) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from prompt-toolkit<3.1.0,>=3.0.30->ipython>=5.3.0->pyvis) (0.2.6)\n",
      "Requirement already satisfied: executing>=1.2.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (1.2.0)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (2.2.1)\n",
      "Requirement already satisfied: pure-eval in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (0.2.2)\n",
      "Requirement already satisfied: six in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from asttokens>=2.1.0->stack-data->ipython>=5.3.0->pyvis) (1.16.0)\n",
      "Installing collected packages: jsonpickle, pyvis\n",
      "Successfully installed jsonpickle-3.0.1 pyvis-0.3.2\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.2.1 is available.\n",
      "You should consider upgrading via the '/Users/loganmarkewich/llama_index/llama-index/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install pyvis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9fe3d26-4f9a-4651-b83f-0018672a34e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kuzugraph_draw.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"kuzugraph_draw.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x8faa728e0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## create graph\n",
    "from pyvis.network import Network\n",
    "\n",
    "g = index.get_networkx_graph()\n",
    "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
    "net.from_nx(g)\n",
    "net.show(\"kuzugraph_draw.html\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "40b97044-d212-4151-bd72-6ea2cff35a29",
   "metadata": {},
   "source": [
    "#### [Optional] Try building the graph and manually add triplets!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9de2ddb-4e82-438b-ba3a-b7680efed944",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SentenceSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "137176d9-1bc2-4203-8379-7b285cd41546",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_parser = SentenceSplitter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc609c08-6fce-444c-84cd-a305fcad6bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = node_parser.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c3ad61-6f2a-4176-96ba-6e9f52d6243d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize an empty database\n",
    "db = kuzu.Database(\"test3\")\n",
    "graph_store = KuzuGraphStore(db)\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
    "index = KnowledgeGraphIndex(\n",
    "    [],\n",
    "    service_context=service_context,\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e03f7e-bb98-4fe0-9fc0-369be2864a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add keyword mappings and nodes manually\n",
    "# add triplets (subject, relationship, object)\n",
    "\n",
    "# for node 0\n",
    "node_0_tups = [\n",
    "    (\"author\", \"worked on\", \"writing\"),\n",
    "    (\"author\", \"worked on\", \"programming\"),\n",
    "]\n",
    "for tup in node_0_tups:\n",
    "    index.upsert_triplet_and_node(tup, nodes[0])\n",
    "\n",
    "# for node 1\n",
    "node_1_tups = [\n",
    "    (\"Interleaf\", \"made software for\", \"creating documents\"),\n",
    "    (\"Interleaf\", \"added\", \"scripting language\"),\n",
    "    (\"software\", \"generate\", \"web sites\"),\n",
    "]\n",
    "for tup in node_1_tups:\n",
    "    index.upsert_triplet_and_node(tup, nodes[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48b1a666-2f84-4524-851a-66efd2beb611",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.knowledge_graph.retriever:> Starting query: Tell me more about Interleaf\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Query keywords: ['Interleaf']\n",
      "ERROR:llama_index.indices.knowledge_graph.retriever:Index was not constructed with embeddings, skipping embedding usage...\n",
      "INFO:llama_index.indices.knowledge_graph.retriever:> Extracted relationships: The following are knowledge sequence in max depth 2 in the form of `subject [predicate, object, predicate_next_hop, object_next_hop ...]`\n",
      "Interleaf ['made software for', 'creating documents']\n",
      "Interleaf ['added', 'scripting language']\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=False, response_mode=\"tree_summarize\"\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about Interleaf\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb4b99d7-452f-4594-94e9-da10a3a23fb8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Interleaf is a software company that specializes in creating documents. They have also added a scripting language to their software.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
